#include <QProgressBar>

#include "../contour/mcconextractor.h"
#include "../contour/molskin.h"
#include "../contour/blurmap.h"
#include "../utils/bilateralfilter.h"
#include "pocketextractor.h"
#include "distancetransform.h"

using namespace PBVP;

PocketExtractor::PocketExtractor(float dist, float isoval, int filter_times)
{
	//m_MarchingDistance = 1.25f;
	//m_IsoValue = 0.01f;
	m_MarchingDistance = dist;
	m_IsoValue = isoval;
	m_FilterTimes = filter_times;

	progressBar = NULL;
}

PocketExtractor::~PocketExtractor()
{
	
}

Reg3DVolumeData* PocketExtractor::computePocketVolume(PBVP::GeometryData *pGeoData)
{
	setProgress(0);
	computeExtent(pGeoData);
	DistanceTransform *pDisTransform = new DistanceTransform(pGeoData);
	pDisTransform->setMinMaxExt(m_MinExt, m_MaxExt);
	setProgress(10);
	// compute SDF for the original surface
	printf("computing SDF for original surface: nv = %d, nt = %d\n", 
		pGeoData->numOfVertices(), pGeoData->numOfTriangles());
	Reg3DVolumeData *pSDFOrig = pDisTransform->computeSDF();
	setProgress(30);
	MCContourExtractor conExtract;
	printf("extracting shell surface from SDF\n");
	GeometryData *shellSurface = conExtract.extractContour(pSDFOrig, m_MarchingDistance);
	printf("Shell surface done: nv = %d, nt = %d\n", 
		shellSurface->numOfVertices(), shellSurface->numOfTriangles());
	setProgress(40);
	delete pDisTransform;
	setProgress(45);
	DistanceTransform *pDisTransform2 = new DistanceTransform(shellSurface);
	pDisTransform2->setMinMaxExt(m_MinExt, m_MaxExt);
	//compute SDF for the shell surface
	printf("compute SDF for shell surface\n");
	Reg3DVolumeData *pSDFShell = pDisTransform2->computeSDF();
	setProgress(70);
	for(unsigned int i = 0; i < pSDFOrig->numOfVerts(); i++ ) {
		double f1 = pSDFOrig->getValue(i);
		double f2 = pSDFShell->getValue(i) - m_MarchingDistance;
		f1 = (f1 < f2)? f1 : f2;
		// if (f1 < 0) f1 = 0;			// set any negative value to 0
		pSDFOrig->setValue(f1, i);
	}
	pSDFOrig->calcMinMax();
	setProgress(80);
	printf("Pocket volume min = %f, max = %f\n", pSDFOrig->getMinValue(), pSDFOrig->getMaxValue());

	if (m_FilterTimes > 0)
	{
		BilateralFilter filter;
		filter.filterIterative(pSDFOrig, m_FilterTimes);
	}
	
	delete pSDFShell;
	delete pDisTransform2;
	return pSDFOrig;
}

Reg3DVolumeData* PocketExtractor::computePocketVolume(const QVector<Atom> &atoms)
{
	MolSkin molSkin;
	GeometryData *pMolSurf = molSkin.constructMolSurface(atoms);
	//BlurMap blurmap;
	//GeometryData *pMolSurf = blurmap.constructSurface(atoms);
	Reg3DVolumeData *pPocketVolume = computePocketVolume(pMolSurf);
	delete pMolSurf;
	return pPocketVolume;
}

GeometryData* PocketExtractor::computePocketSurface(PBVP::GeometryData *pGeoData)
{
	progressBar->show();
	Reg3DVolumeData *pVolume = computePocketVolume(pGeoData);
	MCContourExtractor conExtract;
	printf("extracting pocket surface\n");
	GeometryData *pPocketSurf = conExtract.extractContour(pVolume, m_IsoValue);
	setProgress(90);
	delete pVolume;
	setProgress(100);
	return pPocketSurf;
}

GeometryData* PocketExtractor::computePocketSurface(const QVector<Atom> &atoms)
{
	Reg3DVolumeData *pVolume = computePocketVolume(atoms);
	
	MCContourExtractor conExtract;
	printf("extracting pocket surface\n");
	GeometryData *pPocketSurf = conExtract.extractContour(pVolume, m_IsoValue);
	setProgress(90);
	delete pVolume;
	setProgress(100);
	return pPocketSurf;
}

//////////////////////////////////////////////////////////////////////////
// Private functions
//////////////////////////////////////////////////////////////////////////
void PocketExtractor::computeExtent(PBVP::GeometryData *pGeoData)
{
	m_MinExt = pGeoData->getVertex(0);
	m_MaxExt = pGeoData->getVertex(0);

	for (unsigned int i = 1; i < pGeoData->numOfVertices(); i++)
	{
		Vector3 vert = pGeoData->getVertex(i);
		for (int j = 0; j < 3; j++)
		{
			if (m_MinExt[j] > vert[j]) m_MinExt[j] = vert[j];
			if (m_MaxExt[j] < vert[j]) m_MaxExt[j] = vert[j];
		}
	}

	Vector3 center = (m_MinExt + m_MaxExt) / 2;
	Vector3 half = (m_MaxExt - m_MinExt) / 2;
	half = half + 1.1*m_MarchingDistance;
	m_MinExt = center - half;
	m_MaxExt = center + half;
}

void PocketExtractor::setProgress(int prog)
{
	if(progressBar) {
		progressBar->setValue(prog);
	}
}